# PSCALE VOC数据集处理函数
import torch

def voc_anno_2_coco_target(image_idx:int, anno:dict, class_dict:dict)->dict:
    '''将VOC annotation信息从带层次的dict格式转为coco target的单层形式
        
    '''
    if not "object" in anno:
        raise ValueError('No object found in parameter anno')

    boxes = []
    labels = []
    iscrowd = []
    for obj in anno["object"]:
        xmin = float(obj["bndbox"]["xmin"])
        xmax = float(obj["bndbox"]["xmax"])
        ymin = float(obj["bndbox"]["ymin"])
        ymax = float(obj["bndbox"]["ymax"])

        # 进一步检查数据，有的标注信息中可能有w或h为0的情况，这样的数据会导致计算回归loss为nan
        if xmax <= xmin or ymax <= ymin:
            raise ValueError(f'Error xmax,xmin,ymax,ymin ({xmax},{xmin},{ymax},{ymin})')
            continue
        
        boxes.append([xmin, ymin, xmax, ymax])
        labels.append(class_dict[obj["name"]])
        if "difficult" in obj:
            iscrowd.append(int(obj["difficult"]))
        else:
            iscrowd.append(0)

    # convert everything into a torch.Tensor
    boxes = torch.as_tensor(boxes, dtype=torch.float32)
    labels = torch.as_tensor(labels, dtype=torch.int64)
    iscrowd = torch.as_tensor(iscrowd, dtype=torch.int64)
    image_id = torch.tensor([image_idx])
    area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

    target = {}
    target["boxes"] = boxes
    target["labels"] = labels
    target["image_id"] = image_id
    target["area"] = area
    target["iscrowd"] = iscrowd
    return target
